# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.

import functools
import inspect
import logging


def execute_operation(operation_str, param_dict):
    if operation_str:
        try:
            exec(operation_str, param_dict)
        except Exception as e:
            logging.error(f"Error executing exit code: {e}")


def _get_caller_info(stack):
    """获取调用者的模块名、类名、函数名"""
    caller_module = ""
    caller_class = ""
    caller_function = ""
    try:

        if len(stack) > 1:  # 确保有调用者栈帧
            caller_frame = stack[1].frame
            # 模块名
            caller_module = caller_frame.f_globals.get('__name__', '')
            # 类名（处理self和cls两种情况）
            caller_self = caller_frame.f_locals.get('self') or caller_frame.f_locals.get('cls')
            if caller_self:
                if inspect.isclass(caller_self):
                    caller_class = caller_self.__name__  # 类方法（cls）
                else:
                    caller_class = caller_self.__class__.__name__  # 实例方法（self）
            # 函数名
            caller_function = stack[1].function
    except Exception as e:
        logging.error(f"Error checking direct caller: {e}")
    return caller_module, caller_class, caller_function


def _should_wrap(scope_name, scope_class_name, scope_function, caller_module, caller_class, caller_function):
    """判断是否需要执行包装逻辑（范围匹配）"""
    # 未指定范围时，默认执行包装
    if not all([scope_name, scope_class_name, scope_function]):
        return True
    # 检查调用者是否匹配目标范围
    return (caller_module == scope_name and
            caller_class == scope_class_name and
            caller_function == scope_function)


def marker_prof_wrapper(original_method, params):
    function_name = params.get("function_name", None)
    if function_name is None:
        return original_method

    entry_operation = params.get("entry_operation", None)
    exit_operation = params.get("exit_operation", None)
    scope_name = params.get("scope_name", None)
    scope_class_name = params.get("scope_class_name", None)
    scope_function = params.get("scope_function", None)
    # 策略映射
    strategy_map = {
        True: async_func,
        False: sync_func
    }
    # 判断方法类型，从映射中获取对应的处理策略
    is_async = inspect.iscoroutinefunction(original_method)
    handler = strategy_map[is_async]
    # 执行策略并返回结果
    return handler(original_method, entry_operation, exit_operation, scope_name, scope_class_name,
                   scope_function)


def sync_func(original_method, entry_operation, exit_operation, scope_name, scope_class_name,
              scope_function):
    logging.info(f"<<<INFO: {original_method.__qualname__} is sync function, use sync wrapper")

    @functools.wraps(original_method)
    def wrapper(self, *args, **kwargs):
        entry_args = args
        entry_kwargs = kwargs
        param_dict = {"self": self, "args": entry_args, "kwargs": entry_kwargs}
        # 跑一下获取直接调用函数
        execute_operation(None, {})

        stack = inspect.stack()
        caller_module, caller_class, caller_function = _get_caller_info(stack)

        should_wrap = _should_wrap(
            scope_name, scope_class_name, scope_function,
            caller_module, caller_class, caller_function
        )

        # 如果不需要包装，直接返回原始方法
        if not should_wrap:
            logging.info(
                f"<<<INFO: Skipping wrapper for {original_method.__qualname__} "
                f"due to scope {caller_module}.{caller_class}.{caller_function} mismatch")
            return original_method(self, *args, **kwargs)

        execute_operation(entry_operation, param_dict)

        result = original_method(self, *args, **kwargs)
        param_dict["result"] = result

        execute_operation(exit_operation, param_dict)

        return result

    return wrapper


def async_func(original_method, entry_operation, exit_operation, scope_name, scope_class_name,
               scope_function):
    logging.info(f"<<<INFO: {original_method.__qualname__} is async function, use async wrapper")

    @functools.wraps(original_method)
    async def async_wrapper(self, *args, **kwargs):
        entry_args = args
        entry_kwargs = kwargs
        param_dict = {"self": self, "args": entry_args, "kwargs": entry_kwargs}
        # 跑一下获取直接调用函数
        execute_operation(None, {})

        stack = inspect.stack()

        caller_module, caller_class, caller_function = _get_caller_info(stack)

        should_wrap = _should_wrap(
            scope_name, scope_class_name, scope_function,
            caller_module, caller_class, caller_function
        )

        # 如果不需要包装，直接返回原始方法
        if not should_wrap:
            logging.info(
                f"<<<INFO: Skipping wrapper for {original_method.__qualname__} "
                f"due to scope {caller_module}.{caller_class}.{caller_function} mismatch")
            return original_method(self, *args, **kwargs)

        execute_operation(entry_operation, param_dict)

        result = await original_method(self, *args, **kwargs)
        param_dict["result"] = result

        execute_operation(exit_operation, param_dict)

        return result

    return async_wrapper
